import { GraphModel } from "@tensorflow/tfjs-converter";
import {
  Tensor,
  tensor1d,
  tensor2d,
  tidy,
  image,
  mul,
  sub,
  div,
  squeeze,
  slice,
  sigmoid,
  reshape,
  browser,
  expandDims,
  cast,
  add,
} from "@tensorflow/tfjs-core";
import { Box, createBox, disposeBox } from "./box";
import {
  generateAnchors,
  decodeBounds,
  scaleBoxFromPrediction,
  flipFaceHorizontal,
} from "./utils";

import type {
  Tensor1D,
  Tensor2D,
  Tensor3D,
  Tensor4D,
} from "@tensorflow/tfjs-core";

export type FacePrediction = {
  box: Box;
  landmarks: Tensor2D;
  probability: Tensor1D;
  anchor: Tensor2D | [number, number];
};

export interface NormalizedFace {
  topLeft: [number, number] | Tensor1D;
  bottomRight: [number, number] | Tensor1D;
  landmarks?: number[][] | Tensor2D;
  probability?: number | Tensor1D;
}

const NUM_LANDMARKS = 6;
export class FaceLandmarksDetection {
  private faceModel: GraphModel;
  private width: number;
  private height: number;
  private maxFaces: number;
  private anchors: Tensor2D;
  private anchorsData: number[][];
  private inputSize: Tensor1D;
  private inputSizeData: [number, number];
  private iouThreshold: number;
  private scoreThreshold: number;
  constructor(
    model: GraphModel,
    width: number,
    height: number,
    maxFaces: number,
    iouThreshold: number,
    scoreThreshold: number,
  ) {
    this.faceModel = model;
    this.width = width;
    this.height = height;
    this.maxFaces = maxFaces;
    this.anchorsData = generateAnchors(width, height);
    this.anchors = tensor2d(this.anchorsData);
    this.inputSizeData = [width, height];
    this.inputSize = tensor1d([width, height]);
    this.iouThreshold = iouThreshold;
    this.scoreThreshold = scoreThreshold;
  }
  async getBoundingBoxes(
    inputImage: Tensor4D,
    returnTensors: boolean,
    annotateBoxes = true,
  ): Promise<{
    boxes: Array<FacePrediction | Box>;
    scaleFactor: Tensor | [number, number];
  }> {
    const [detectedOutputs, boxes, scores] = tidy(
      (): [Tensor2D, Tensor2D, Tensor1D] => {
        const resizedImage = image.resizeBilinear(inputImage, [
          this.width,
          this.height,
        ]);
        const normalizedImage = mul(sub(div(resizedImage, 255), 0.5), 2);

        const batchedPrediction = this.faceModel.predict(normalizedImage);
        const prediction = squeeze(batchedPrediction as Tensor3D);

        const decodedBounds = decodeBounds(
          prediction as Tensor2D,
          this.anchors,
          this.inputSize,
        );
        const logits = slice(prediction as Tensor2D, [0, 0], [-1, 1]);
        const scores = squeeze(sigmoid(logits));
        return [prediction as Tensor2D, decodedBounds, scores as Tensor1D];
      },
    );

    const boxIndicesTensor = image.nonMaxSuppression(
      boxes,
      scores,
      this.maxFaces,
      this.iouThreshold,
      this.scoreThreshold,
    );

    const boxIndices = await boxIndicesTensor.array();
    boxIndicesTensor.dispose();

    let boundingBoxes: Tensor[] | number[][][] = boxIndices.map(
      (boxIndex: number) => slice(boxes, [boxIndex, 0], [1, -1]),
    );
    if (!returnTensors) {
      boundingBoxes = await Promise.all(
        boundingBoxes.map(async (boundingBox: any) => {
          const vals = await boundingBox.array();
          boundingBox.dispose();
          return vals;
        }),
      );
    }

    const originalHeight = inputImage.shape[1];
    const originalWidth = inputImage.shape[2];

    let scaleFactor: Tensor | [number, number];
    if (returnTensors) {
      scaleFactor = div([originalWidth, originalHeight], this.inputSize);
    } else {
      scaleFactor = [
        originalWidth / this.inputSizeData[0],
        originalHeight / this.inputSizeData[1],
      ];
    }

    const annotatedBoxes = [];
    for (let i = 0; i < boundingBoxes.length; i++) {
      const boundingBox = boundingBoxes[i] as Tensor2D | number[][];
      const annotatedBox = tidy(() => {
        const box =
          boundingBox instanceof Tensor
            ? createBox(boundingBox)
            : createBox(tensor2d(boundingBox));

        if (!annotateBoxes) {
          return box;
        }

        const boxIndex = boxIndices[i];

        let anchor;
        if (returnTensors) {
          anchor = slice(this.anchors, [boxIndex, 0], [1, 2]);
        } else {
          anchor = this.anchorsData[boxIndex] as [number, number];
        }

        const landmarks = reshape(
          squeeze(
            slice(detectedOutputs, [boxIndex, NUM_LANDMARKS - 1], [1, -1]),
          ),
          [NUM_LANDMARKS, -1],
        );
        const probability = slice(scores, [boxIndex], [1]);

        return { box, landmarks, probability, anchor };
      });
      annotatedBoxes.push(annotatedBox);
    }

    boxes.dispose();
    scores.dispose();
    detectedOutputs.dispose();

    return {
      boxes: annotatedBoxes as Array<FacePrediction | Box>,
      scaleFactor,
    };
  }
  async estimateFaces(
    input:
      | Tensor3D
      | ImageData
      | HTMLVideoElement
      | HTMLImageElement
      | HTMLCanvasElement,
    returnTensors = false,
    flipHorizontal = false,
    annotateBoxes = true,
  ): Promise<NormalizedFace[]> {
    const [, width] =
      input instanceof Tensor
        ? [input.shape[0], input.shape[1]]
        : [input.height, input.width];
    const image = tidy(() => {
      if (!(input instanceof Tensor)) {
        input = browser.fromPixels(input);
      }
      return expandDims(cast(input as Tensor, "float32"), 0);
    });
    const { boxes, scaleFactor } = await this.getBoundingBoxes(
      image as Tensor4D,
      returnTensors,
      annotateBoxes,
    );
    image.dispose();

    if (returnTensors) {
      return boxes.map((face: FacePrediction | Box) => {
        const scaledBox = scaleBoxFromPrediction(face, scaleFactor as Tensor1D);
        let normalizedFace: NormalizedFace = {
          topLeft: slice(scaledBox, [0], [2]) as Tensor1D,
          bottomRight: slice(scaledBox, [2], [2]) as Tensor1D,
        };

        if (annotateBoxes) {
          const { landmarks, probability, anchor } = face as {
            landmarks: Tensor2D;
            probability: Tensor1D;
            anchor: Tensor2D | [number, number];
          };

          const normalizedLandmarks: Tensor2D = mul(
            add(landmarks, anchor),
            scaleFactor,
          );
          normalizedFace.landmarks = normalizedLandmarks;
          normalizedFace.probability = probability;
        }

        if (flipHorizontal) {
          normalizedFace = flipFaceHorizontal(normalizedFace, width);
        }
        return normalizedFace;
      });
    }

    return Promise.all(
      // @ts-ignore
      boxes.map(async (face: FacePrediction) => {
        const scaledBox = scaleBoxFromPrediction(
          face,
          scaleFactor as [number, number],
        );
        let normalizedFace: NormalizedFace;
        if (!annotateBoxes) {
          const boxData = await scaledBox.array();
          normalizedFace = {
            topLeft: (boxData as number[]).slice(0, 2) as [number, number],
            bottomRight: (boxData as number[]).slice(2) as [number, number],
          };
        } else {
          const [landmarkData, boxData, probabilityData] = await Promise.all(
            [face.landmarks, scaledBox, face.probability].map(async (d) =>
              d.array(),
            ),
          );

          const anchor = face.anchor as [number, number];
          const [scaleFactorX, scaleFactorY] = scaleFactor as [number, number];
          const scaledLandmarks = (landmarkData as Array<[number, number]>).map(
            (landmark) => [
              (landmark[0] + anchor[0]) * scaleFactorX,
              (landmark[1] + anchor[1]) * scaleFactorY,
            ],
          );

          normalizedFace = {
            topLeft: (boxData as number[]).slice(0, 2) as [number, number],
            bottomRight: (boxData as number[]).slice(2) as [number, number],
            landmarks: scaledLandmarks,
            probability: probabilityData as number,
          };

          disposeBox(face.box);
          face.landmarks.dispose();
          face.probability.dispose();
        }

        scaledBox.dispose();

        if (flipHorizontal) {
          normalizedFace = flipFaceHorizontal(normalizedFace, width);
        }

        return normalizedFace;
      }),
    );
  }
}

export default FaceLandmarksDetection;
